//
// Created by 朱林 on 2020/9/21.
//

#ifndef PERSONPOSE_MULTI_POSE_H
#define PERSONPOSE_MULTI_POSE_H

#include <vector>
#include "opencv2/opencv.hpp"
#include <math.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <MNN/ImageProcess.hpp>
#include <MNN/Interpreter.hpp>
#include <MNN/AutoTime.hpp>

#define MNN_OPEN_TIME_TRACE
#define OUTPUT_STRIDE 16
#define MAX_POSE_DETECTIONS 10
#define NUM_KEYPOINTS 17
#define SCORE_THRESHOLD 0.5
#define MIN_POSE_SCORE 0.25
#define NMS_RADIUS 20
#define LOCAL_MAXIMUM_RADIUS 1
#define OFFSET_NODE_NAME "offset_2"
#define DISPLACE_FWD_NODE_NAME "displacement_fwd_2"
#define DISPLACE_BWD_NODE_NAME "displacement_bwd_2"
#define HEATMAPS "heatmap"
#define CIRCLE_RADIUS 3

using namespace MNN;


typedef struct Point {
    float x;
    float y;
} Point;

const std::vector<std::string> PoseNames{"nose", "leftEye", "rightEye", "leftEar", "rightEar",
                                         "leftShoulder", "rightShoulder", "leftElbow", "rightElbow",
                                         "leftWrist",
                                         "rightWrist", "leftHip", "rightHip", "leftKnee",
                                         "rightKnee",
                                         "leftAnkle", "rightAnkle"};

const std::vector<std::pair<std::string, std::string>> PoseChain{
        {"nose",          "leftEye"},
        {"leftEye",       "leftEar"},
        {"nose",          "rightEye"},
        {"rightEye",      "rightEar"},
        {"nose",          "leftShoulder"},
        {"leftShoulder",  "leftElbow"},
        {"leftElbow",     "leftWrist"},
        {"leftShoulder",  "leftHip"},
        {"leftHip",       "leftKnee"},
        {"leftKnee",      "leftAnkle"},
        {"nose",          "rightShoulder"},
        {"rightShoulder", "rightElbow"},
        {"rightElbow",    "rightWrist"},
        {"rightShoulder", "rightHip"},
        {"rightHip",      "rightKnee"},
        {"rightKnee",     "rightAnkle"}};

class multiPose {
public:
    multiPose();

    ~multiPose();

    bool init(const std::string &mnn_path, int num_thread_ = 4, bool openCL = true);

    std::vector<Point> detect(cv::Mat &img);

private:
    std::shared_ptr<MNN::Interpreter> pose_interpreter;

    MNN::Session *pose_session = nullptr;

    MNN::Tensor *input_tensor = nullptr;

    int num_thread = 4;

    const float means[3] = {127.5f, 127.5f, 127.5f};

    const float norms[3] = {2.0f / 255.0f, 2.0f / 255.0f, 2.0f / 255.0f};

};


#endif //PERSONPOSE_MULTI_POSE_H
